use std::collections::HashSet;

type Coord = (i32, i32);
type Image = HashSet<Coord>;

fn parse_input() -> (Vec<bool>, Image) {
    let input = std::fs::read_to_string("./day20/input.txt").unwrap();
    let algorithm = input.lines().nth(0).unwrap().chars().map(|ch| ch == '#').collect();

    let mut image = Image::new();
    for (y, line) in input.lines().skip(2).enumerate() {
        for (x, ch) in line.chars().enumerate() {
            if ch == '#' {
                image.insert((x as i32, y as i32));
            }
        }
    }

    (algorithm, image)
}

fn adjacent_coords(&(x ,y): &Coord) -> [Coord; 9] {
    [
        (x - 1, y - 1), (x    , y - 1), (x + 1, y - 1),
        (x - 1, y    ), (x    , y    ), (x + 1, y    ),
        (x - 1, y + 1), (x    , y + 1), (x + 1, y + 1),
    ]
}

fn advance_image(image: Image, outside_algo: bool ,algorithm: &Vec<bool>) -> (Image, bool) {
    let &min_x = image.iter().map(|(x, _)| x).min().unwrap();
    let &max_x = image.iter().map(|(x, _)| x).max().unwrap();
    let &min_y = image.iter().map(|(_, y)| y).min().unwrap();
    let &max_y = image.iter().map(|(_, y)| y).max().unwrap();

    let out_range = |&(x, y): &Coord| -> bool {
        !(min_x <= x && x <= max_y && min_y <= y && y <= max_y)
    };

    let mut new_image = Image::new();

    for y in min_y-1..=max_y+1 {
        for x in min_x-1..=max_x+1 {
            let coord = (x, y);
            let bin_str: String = adjacent_coords(&coord).into_iter().map(|c| -> char {
                if image.contains(&c) || (outside_algo && out_range(&c)) { 
                    '1' 
                } else { 
                    '0' 
                }
            }).collect();
            
            let index = usize::from_str_radix(&bin_str, 2).unwrap();
            if algorithm[index] {
                new_image.insert(coord);
            }
        }
    }

    (new_image, if outside_algo { algorithm[511] } else { algorithm[0] })
}

#[allow(unused)]
fn part1(image: Image, algorithm: &Vec<bool>) {
    let image = (0..2).into_iter().fold((image, false), |(image, outside_algo), _| {
        advance_image(image, outside_algo, algorithm)
    }).0;
    println!("Part1: {}", image.len());
}

#[allow(unused)]
fn part2(image: Image, algorithm: &Vec<bool>) {
    let image = (0..50).into_iter().fold((image, false), |(image, outside_algo), _| {
        advance_image(image, outside_algo, algorithm)
    }).0;
    println!("Part2: {}", image.len());
}

fn main() {
    let (algorithm, image) = parse_input();
    part1(image.clone(), &algorithm);
    part2(image, &algorithm);
}
